import sqlite3
from typing import Iterable
import pandas as pd

class Database:
    def __init__(self, db_path):
        self.db_path = db_path
        self.conn = sqlite3.connect(self.db_path)
        self.conn.execute("PRAGMA foreign_keys = ON;")
        self.create_tables()

    def create_tables(self):
        cur = self.conn.cursor()
        # 开户信息表（标准列对齐 normalizer）
        cur.execute("""
        CREATE TABLE IF NOT EXISTS accounts (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            账号 TEXT, 卡号 TEXT, 开户名 TEXT, 开户行名称 TEXT, 证件号码 TEXT, 状态 TEXT, 余额 REAL,
            销户日期 TEXT, 销户前余额 REAL, 手机号 TEXT, 证件类型 TEXT, 账号类型 TEXT, 币种 TEXT,
            开户日期 TEXT, 开户网点 TEXT
        );
        """)
        # 银行流水表（标准列对齐 normalizer）
        cur.execute("""
        CREATE TABLE IF NOT EXISTS transactions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            交易时间 TEXT, 交易时间戳 TEXT, 交易流水号 TEXT, 交易账号 TEXT, 户名 TEXT,
            摘要用途 TEXT, 交易类型 TEXT, 金额 REAL, 币种 TEXT, 余额 REAL,
            对方户名 TEXT, 对方账号 TEXT, 对方开户行 TEXT, 对方证件号 TEXT,
            交易渠道 TEXT, 交易网点 TEXT, 备注 TEXT, 备注2 TEXT, 交易状态 TEXT
        );
        """)
        self.conn.commit()

    def _insert_dataframe(self, table: str, df: pd.DataFrame):
        if df is None or df.empty:
            return 0
        # 列名可能包含'摘要/用途'，统一一下兼容
        cols = [c.replace('/', '') for c in df.columns]  # 将 摘要/用途 -> 摘要用途
        df2 = df.copy()
        df2.columns = cols
        df2 = df2.where(pd.notnull(df2), None)
        placeholders = ",".join(["?"] * len(cols))
        col_sql = ",".join([f'"{c}"' for c in cols])
        cur = self.conn.cursor()
        cur.executemany(
            f"INSERT INTO {table} ({col_sql}) VALUES ({placeholders});",
            df2.values.tolist()
        )
        self.conn.commit()
        return len(df2)

    def insert_accounts(self, df: pd.DataFrame) -> int:
        return self._insert_dataframe("accounts", df)

    def insert_transactions(self, df: pd.DataFrame) -> int:
        return self._insert_dataframe("transactions", df)

    def close(self):
        self.conn.close()
